package top.tenebrous.termbor.rpc.registry.zookeeper;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import top.tenebrous.termbor.rpc.common.helper.RpcServiceHelper;
import top.tenebrous.termbor.rpc.loadbalance.ServiceLoadBalance;
import top.tenebrous.termbor.rpc.loadbalance.random.RandomServiceLoadBalance;
import top.tenebrous.termbor.rpc.protocol.meta.ServiceMeta;
import top.tenebrous.termbor.rpc.registry.RegistryService;
import top.tenebrous.termbor.rpc.registry.config.RegistryConfig;
import top.tenebrous.termbor.rpc.spi.SPIClass;
import top.tenebrous.termbor.rpc.spi.loader.ExtensionLoader;

import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.Random;

/**
 * @program: termbor-rpc
 * @description:
 * @author: yazhang6
 * @create: 2024-05-23 09:36
 **/
@SPIClass
public class ZookeeperRegistryService implements RegistryService {
    public static final int BASE_SLEEP_TIME_MS = 1000;
    public static final int MAX_RETRIES = 3;
    public static final String ZK_BASE_PATH = "/rpc";
    private ServiceDiscovery<ServiceMeta> serviceDiscovery;

    private ServiceLoadBalance<ServiceInstance<ServiceMeta>> serviceLoadBalance;

    @Override
    public void init(RegistryConfig registryConfig) throws Exception {
        CuratorFramework client = CuratorFrameworkFactory.newClient(registryConfig.getRegistryAddr(), new ExponentialBackoffRetry(BASE_SLEEP_TIME_MS, MAX_RETRIES));
        client.start();
        JsonInstanceSerializer<ServiceMeta> serializer = new JsonInstanceSerializer<>(ServiceMeta.class);
        this.serviceDiscovery = ServiceDiscoveryBuilder.builder(ServiceMeta.class)
            .client(client)
            .serializer(serializer)
            .basePath(ZK_BASE_PATH)
            .build();
        this.serviceDiscovery.start();
        this.serviceLoadBalance = ExtensionLoader.getExtension(ServiceLoadBalance.class, registryConfig.getRegistryLoadBalanceType());
    }

    @Override
    public void register(ServiceMeta serviceMeta) throws Exception {
        ServiceInstance<ServiceMeta> serviceInstance = ServiceInstance
            .<ServiceMeta>builder()
            .name(RpcServiceHelper.buildServiceKey(serviceMeta.getServiceName(), serviceMeta.getServiceVersion(), serviceMeta.getServiceGroup()))
            .address(serviceMeta.getServiceAddr())
            .port(serviceMeta.getServicePort())
            .payload(serviceMeta)
            .build();
        serviceDiscovery.registerService(serviceInstance);

    }

    @Override
    public void unRegister(ServiceMeta serviceMeta) throws Exception {
        ServiceInstance<ServiceMeta> serviceInstance =
            ServiceInstance.<ServiceMeta>builder().name(serviceMeta.getServiceName())
                .address(serviceMeta.getServiceAddr()).port(serviceMeta.getServicePort()).payload(serviceMeta).build();
        serviceDiscovery.unregisterService(serviceInstance);
    }

    @Override
    public ServiceMeta discovery(String serviceName, int invokerHashCode) throws Exception {
        Collection<ServiceInstance<ServiceMeta>> serviceInstances = serviceDiscovery.queryForInstances(serviceName);
        ServiceInstance<ServiceMeta> instance = serviceLoadBalance.select((List<ServiceInstance<ServiceMeta>>)serviceInstances, invokerHashCode);
        if (instance != null) {
            System.out.println("注册中心的服务数据: " + instance.getPayload());
            return instance.getPayload();
        }
        return null;
    }

    @Override
    public void destroy() throws IOException {
        if (serviceDiscovery != null) {
            serviceDiscovery.close();
        }
    }
}
